#include <chrono>
#include <cstddef>
#include <iostream>
#include <random>
#include <vector>
#include <omp.h>
#include <cmath>
#include <immintrin.h>  // 引入AVX2指令集头文件

using std::vector;

class FigureProcessor {
//private:
//    vector<vector<unsigned char>> figure;
//    vector<vector<unsigned char>> result;
//    const size_t size;
private:
    // 改为一维数组
    std::vector<unsigned char> figure; 
    std::vector<unsigned char> result;
    const size_t size;

    // 获取一维数组的索引
    inline size_t idx(size_t row, size_t col) const {
        return row * size + col;
    }


public:
    //FigureProcessor(size_t size, size_t seed = 0) : size(size) {

    //    // !!! Please do not modify the following code !!!
    //    // 如果你需要修改内存的数据结构，请不要修改初始化的顺序和逻辑
    //    // 助教可能会通过指定某个初始化seed 的方式来验证你的代码
    //    // 如果你修改了初始化的顺序，可能会导致你的代码无法通过测试
    //    std::random_device rd;
    //    std::mt19937_64 gen(seed == 0 ? rd() : seed);
    //    std::uniform_int_distribution<unsigned char> distribution(0, 255);
    //    // !!! ----------------------------------------- !!!

    //    // 两个数组的初始化在这里，可以改动，但请注意 gen 的顺序是从上到下从左到右即可。
    //    for (size_t i = 0; i < size; ++i) {
    //        figure.push_back(vector<unsigned char>());
    //        for (size_t j = 0; j < size; ++j) {
    //            figure[i].push_back(static_cast<unsigned char>(distribution(gen)));
    //        }
    //    }

    //    for (size_t i = 0; i < size; ++i) {
    //        result.push_back(vector<unsigned char>());
    //        for (size_t j = 0; j < size; ++j) {
    //            result[i].push_back(0);
    //        }
    //    }
    //}
    FigureProcessor(size_t size, size_t seed = 0) : size(size), figure(size* size), result(size* size) {
        // 初始化随机数生成器（不变）
        std::random_device rd;
        std::mt19937_64 gen(seed == 0 ? rd() : seed);
        std::uniform_int_distribution<unsigned char> distribution(0, 255);

        // 填充 figure 数据
        for (size_t i = 0; i < size; ++i) {
            for (size_t j = 0; j < size; ++j) {
                figure[idx(i, j)] = static_cast<unsigned char>(distribution(gen));
            }
        }

        // 初始化 result 数据为 0
        std::fill(result.begin(), result.end(), 0);
    }


    ~FigureProcessor() = default;

    // Gaussian filter
    // [[1, 2, 1], [2, 4, 2], [1, 2, 1]] / 16
    //FIXME: Feel free to optimize this function
    //Hint: You can use SIMD instructions to optimize this function
    void gaussianFilter() {
        // 定义高斯权重 (3x3 的矩阵)
        const int weights[3][3] = {
            {1, 2, 1},
            {2, 4, 2},
            {1, 2, 1}
        };

        //处理内部区域
#pragma omp parallel for
        for (size_t i = 1; i < size - 1; ++i) { // 遍历行，跳过边界
            for (size_t j = 1; j < size - 1; j += 8) { // 每次处理8个像素
                // 用于累加高斯权重结果
                __m256i sum = _mm256_setzero_si256();

                // 遍历高斯滤波器的3x3区域
                for (int ki = -1; ki <= 1; ++ki) {
                    for (int kj = -1; kj <= 1; ++kj) {
                        // 加载8个像素
                        __m256i pixels = _mm256_cvtepu8_epi16(
                            _mm_loadl_epi64(reinterpret_cast<const __m128i*>(&figure[idx(i + ki, j + kj)]))
                        );

                        // 乘以权重
                        __m256i weight = _mm256_set1_epi16(weights[ki + 1][kj + 1]);
                        __m256i weighted = _mm256_mullo_epi16(pixels, weight);

                        // 累加结果
                        sum = _mm256_add_epi16(sum, weighted);
                    }
                }

                // 将结果除以权重总和(16)并截断为8位
                __m256i resultVec = _mm256_srli_epi16(sum, 4); // 右移4位，除以16
                __m256i result8 = _mm256_packus_epi16(resultVec, resultVec);

                // 存储结果
                _mm_storel_epi64(reinterpret_cast<__m128i*>(&result[idx(i, j)]), _mm256_castsi256_si128(result8));
            }
        }
        // 处理边界和角点
#pragma omp for
        for (size_t i = 1; i < size - 1; ++i) {
            result[idx(i, 0)] =
                (figure[idx(i - 1, 0)] + 2 * figure[idx(i - 1, 0)] + figure[idx(i - 1, 1)] +
                    2 * figure[idx(i, 0)] + 4 * figure[idx(i, 0)] + 2 * figure[idx(i, 1)] +
                    figure[idx(i + 1, 0)] + 2 * figure[idx(i + 1, 0)] + figure[idx(i + 1, 1)]) /
                16;

            result[idx(i, size - 1)] =
                (figure[idx(i - 1, size - 2)] + 2 * figure[idx(i - 1, size - 1)] +
                    figure[idx(i - 1, size - 1)] + 2 * figure[idx(i, size - 2)] +
                    4 * figure[idx(i, size - 1)] + 2 * figure[idx(i, size - 1)] +
                    figure[idx(i + 1, size - 2)] + 2 * figure[idx(i + 1, size - 1)] +
                    figure[idx(i + 1, size - 1)]) /
                16;
        }

#pragma omp for
        for (size_t j = 1; j < size - 1; ++j) {
            result[idx(0, j)] =
                (figure[idx(0, j - 1)] + 2 * figure[idx(0, j)] + figure[idx(0, j + 1)] +
                    2 * figure[idx(0, j - 1)] + 4 * figure[idx(0, j)] + 2 * figure[idx(0, j + 1)] +
                    figure[idx(1, j - 1)] + 2 * figure[idx(1, j)] + figure[idx(1, j + 1)]) /
                16;

            result[idx(size - 1, j)] =
                (figure[idx(size - 2, j - 1)] + 2 * figure[idx(size - 2, j)] +
                    figure[idx(size - 2, j + 1)] + 2 * figure[idx(size - 1, j - 1)] +
                    4 * figure[idx(size - 1, j)] + 2 * figure[idx(size - 1, j + 1)] +
                    figure[idx(size - 1, j - 1)] + 2 * figure[idx(size - 1, j)] +
                    figure[idx(size - 1, j + 1)]) /
                16;
        }

        // 处理四个角点
        // 左上角
        result[idx(0, 0)] = (4 * figure[idx(0, 0)] + 2 * figure[idx(0, 1)] + 2 * figure[idx(1, 0)] +
            figure[idx(1, 1)]) /
            9;

        // 右上角
        result[idx(0, size - 1)] = (4 * figure[idx(0, size - 1)] + 2 * figure[idx(0, size - 2)] +
            2 * figure[idx(1, size - 1)] + figure[idx(1, size - 2)]) /
            9;

        // 左下角
        result[idx(size - 1, 0)] = (4 * figure[idx(size - 1, 0)] + 2 * figure[idx(size - 1, 1)] +
            2 * figure[idx(size - 2, 0)] + figure[idx(size - 2, 1)]) /
            9;

        // 右下角
        result[idx(size - 1, size - 1)] =
            (4 * figure[idx(size - 1, size - 1)] + 2 * figure[idx(size - 1, size - 2)] +
                2 * figure[idx(size - 2, size - 1)] + figure[idx(size - 2, size - 2)]) /
            9;
    }

    // Power law transformation
    // FIXME: Feel free to optimize this function
    // Hint: LUT to optimize this function?
    void powerLawTransformation() {
        // 创建查找表(LUT)
        unsigned char LUT[256];
        constexpr float gamma = 0.5f;
        const float delta = std::pow(255.0f, gamma);

        // 预计算查找表
        for (size_t i = 0; i < 256; ++i) {
            LUT[i] = static_cast<unsigned char>(255.0f * std::pow(i, gamma) / delta + 0.5f);
        }

#pragma omp parallel for
        for (size_t i = 0; i < size; ++i) {
            for (size_t j = 0; j < size; j += 32) { // 每次处理32字节
                // 加载32字节像素
                __m256i pixels = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&figure[idx(i, j)]));

                // 将像素分成低16字节和高16字节
                __m128i low = _mm256_castsi256_si128(pixels);
                __m128i high = _mm256_extracti128_si256(pixels, 1);

                // 逐字节查找
                __m128i resultLow = _mm_set_epi8(
                    LUT[_mm_extract_epi8(low, 15)], LUT[_mm_extract_epi8(low, 14)],
                    LUT[_mm_extract_epi8(low, 13)], LUT[_mm_extract_epi8(low, 12)],
                    LUT[_mm_extract_epi8(low, 11)], LUT[_mm_extract_epi8(low, 10)],
                    LUT[_mm_extract_epi8(low, 9)], LUT[_mm_extract_epi8(low, 8)],
                    LUT[_mm_extract_epi8(low, 7)], LUT[_mm_extract_epi8(low, 6)],
                    LUT[_mm_extract_epi8(low, 5)], LUT[_mm_extract_epi8(low, 4)],
                    LUT[_mm_extract_epi8(low, 3)], LUT[_mm_extract_epi8(low, 2)],
                    LUT[_mm_extract_epi8(low, 1)], LUT[_mm_extract_epi8(low, 0)]
                );

                __m128i resultHigh = _mm_set_epi8(
                    LUT[_mm_extract_epi8(high, 15)], LUT[_mm_extract_epi8(high, 14)],
                    LUT[_mm_extract_epi8(high, 13)], LUT[_mm_extract_epi8(high, 12)],
                    LUT[_mm_extract_epi8(high, 11)], LUT[_mm_extract_epi8(high, 10)],
                    LUT[_mm_extract_epi8(high, 9)], LUT[_mm_extract_epi8(high, 8)],
                    LUT[_mm_extract_epi8(high, 7)], LUT[_mm_extract_epi8(high, 6)],
                    LUT[_mm_extract_epi8(high, 5)], LUT[_mm_extract_epi8(high, 4)],
                    LUT[_mm_extract_epi8(high, 3)], LUT[_mm_extract_epi8(high, 2)],
                    LUT[_mm_extract_epi8(high, 1)], LUT[_mm_extract_epi8(high, 0)]
                );

                // 合并低高16字节并存储结果
                __m256i resultVec = _mm256_set_m128i(resultHigh, resultLow);
                _mm256_storeu_si256(reinterpret_cast<__m256i*>(&result[idx(i, j)]), resultVec);
            }
        }
    }

    // Run benchmark
    unsigned int calcChecksum() {
        unsigned int sum = 0;
        constexpr size_t mod = 1000000007;
        for (size_t i = 0; i < size; ++i) {
            for (size_t j = 0; j < size; ++j) {
                sum += result[idx(i, j)]; // 改为一维数组
                sum %= mod;
            }
        }
        return sum;
    }
    void runBenchmark() {
        auto start = std::chrono::high_resolution_clock::now();
        gaussianFilter();
        auto middle = std::chrono::high_resolution_clock::now();

        unsigned int sum = calcChecksum();

        auto middle2 = std::chrono::high_resolution_clock::now();
        powerLawTransformation();
        auto end = std::chrono::high_resolution_clock::now();

        sum += calcChecksum();
        sum %= 1000000007;
        std::cout << "Checksum: " << sum << "\n";

        auto milliseconds =
            std::chrono::duration_cast<std::chrono::milliseconds>(middle - start) +
            std::chrono::duration_cast<std::chrono::milliseconds>(end - middle2);
        std::cout << "Benchmark time: " << milliseconds.count() << " ms\n";
    }
};

// Main function
// !!! Please do not modify the main function !!!
int main(int argc, const char** argv) {
    omp_set_num_threads(4);
    constexpr size_t size = 16384;
    FigureProcessor processor(size, argc > 1 ? std::stoul(argv[1]) : 0);
    processor.runBenchmark();
    return 0;
}
